
import os

import dotenv
from langchain_classic.chains.llm import LLMChain
from langchain_classic.chains.sequential import SimpleSequentialChain, SequentialChain
from langchain_classic.chains.sql_database.query import create_sql_query_chain
from langchain_community.utilities import SQLDatabase
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
from langchain_openai import ChatOpenAI



# LLMChain的使用,LLMChain在1.0版本中被弃用，请使用ChatChain代替，这里只做演示,使用PromptTemplate
def fun_1():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )

    chain = LLMChain(
        llm=chat_model
        , prompt=PromptTemplate.from_template("{input}")
    )
    response = chain.invoke({"input": "介绍下LLM"})
    print(response)

#fun_1()

#LLMChain的使用,LLMChain在1.0版本中被弃用，请使用ChatChain代替，这里只做演示,使用ChatPromptTemplate
def fun_2():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )

    chain = LLMChain(
        llm=chat_model
        , prompt=ChatPromptTemplate.from_messages(["human", "{input}"])
        ,verbose= True # 输出详细日志
    )
    response = chain.invoke({"input": "介绍下LLM"})
    print(response)

#fun_2()


#SimpleSequentialChain的使用,最简单的顺序链，顺序执行多个链，每个步骤都有单一的输入和输出。
def fun_3():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )
    #注意，因为是SimpleSequentialChain的使用，输入必须是是input
    chainA_template = ChatPromptTemplate.from_messages(
        [("system" , "你是一个有礼貌的AI") , ("human" , "{input}")]
    )
    chainA = LLMChain(
        llm=chat_model
        , prompt=chainA_template
        , verbose= True
    )

    #description没有显示的跟上一个链关联，因为是SimpleSequentialChain的使用，上一个链的输出作为下一个链的输入
    chainB_template = ChatPromptTemplate.from_messages(
        [("system", "你是一个话少的程序员"), ("human", "针对上一个提问的完整说明:{description}，请给出一个代码示例，并给出注释。")]
    )
    chainB = LLMChain(
        llm=chat_model
        , prompt=chainB_template
        , verbose=True
    )

    full_chain = SimpleSequentialChain(
        chains=[chainA, chainB]
        , verbose=True
    )
    response = full_chain.invoke(input={"input": "介绍下PYTHON"})
    print(response)

#fun_3()

#SequentialChain,允许不同子链有独立的输入输出，支持分支、条件罗技。
def fun_4():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )
    #A
    chainA_template = ChatPromptTemplate.from_messages(
        [("system", "你是一个精通各知识领域的AI"), ("human", "请你先尽可能详细的解释下{knowledge},并且{action}")]
    )
    chainA = LLMChain(
        llm=chat_model
        , prompt=chainA_template
        , verbose=True
        , output_key="output_a"
    )
    #B
    chainB_template = ChatPromptTemplate.from_messages(
        [
            ("system", "你非常善于提取文本中重要信息，并做出简短总结")
            , ("human", "这是针对第一个提问的完整解释内容:{output_a}")
            , ("human", "请你根据上述说明，尽可能简短输出重要结论，控制在100字内")
         ]
    )
    chainB = LLMChain(
        llm=chat_model
        , prompt=chainB_template
        , verbose=True
        , output_key="output_b"
    )

    full_chain = SequentialChain(
        chains=[chainA, chainB]
        , verbose=True
        , input_variables=["knowledge", "action"]
        , output_variables=["output_a" , "output_b"] #把第一个模型和第二个模型的输出都输出
    )

    response = full_chain.invoke(input={"knowledge": "中国足球为什么踢得这么烂" , "action": "举个例子"})
    print(response)

#fun_4()

#create_sql_query_chain  (pip install pymysql),将自然语言转换成SQL查询
def fun_5():
    db = SQLDatabase.from_uri("mysql+pymysql://root:yaoyiwen@127.0.0.1:3306/test")
    #print(db.get_usable_table_names()) #测试连接，打印表名
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )

    chain = create_sql_query_chain(llm=chat_model , db=db)
    response = chain.invoke({"question": "请查询sys_user表里所有的用户名,不限制返回的数量"})
    print(response)
fun_5()
